{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num_vectors: 1000000 dim: 128\n",
      "Reading...\n",
      "Read done\n",
      "Training...\n",
      "[1] Setting NumberOfThreads with value 4\n",
      "Training Done...\n",
      "[1] Setting Samples with value 1000\n",
      "[1] Start to build KDTree 1\n",
      "[1] 1 KDTree built, 999998 1000000\n",
      "[1] Build Tree time (s): 2\n",
      "[1] build RNG graph!\n",
      "[1] Parallel TpTree Partition begin\n",
      "[1] Finish Getting Leaves for Tree 0\n",
      "[1] Finish Getting Leaves for Tree 1\n",
      "[1] Finish Getting Leaves for Tree 3\n",
      "[1] Finish Getting Leaves for Tree 2\n",
      "[1] Finish Getting Leaves for Tree 4\n",
      "[1] Finish Getting Leaves for Tree 5\n",
      "[1] Finish Getting Leaves for Tree 6\n",
      "[1] Finish Getting Leaves for Tree 7\n",
      "[1] Finish Getting Leaves for Tree 8\n",
      "[1] Finish Getting Leaves for Tree 9\n",
      "[1] Finish Getting Leaves for Tree 10\n",
      "[1] Finish Getting Leaves for Tree 11\n",
      "[1] Finish Getting Leaves for Tree 12\n",
      "[1] Finish Getting Leaves for Tree 13\n",
      "[1] Finish Getting Leaves for Tree 14\n",
      "[1] Finish Getting Leaves for Tree 15\n",
      "[1] Finish Getting Leaves for Tree 16\n",
      "[1] Finish Getting Leaves for Tree 17\n",
      "[1] Finish Getting Leaves for Tree 18\n",
      "[1] Finish Getting Leaves for Tree 19\n",
      "[1] Finish Getting Leaves for Tree 20\n",
      "[1] Finish Getting Leaves for索引构建完成并已保存。\n",
      " Tree 21\n",
      "[1] Finish Getting Leaves for Tree 22\n",
      "[1] Finish Getting Leaves for Tree 23\n",
      "[1] Finish Getting Leaves for Tree 24\n",
      "[1] Finish Getting Leaves for Tree 25\n",
      "[1] Finish Getting Leaves for Tree 26\n",
      "[1] Finish Getting Leaves for Tree 27\n",
      "[1] Finish Getting Leaves for Tree 28\n",
      "[1] Finish Getting Leaves for Tree 29\n",
      "[1] Finish Getting Leaves for Tree 30\n",
      "[1] Finish Getting Leaves for Tree 31\n",
      "[1] Parallel TpTree Partition done\n",
      "[1] Build TPTree time (s): 31\n",
      "[1] Processing Tree 0 0%\n",
      "[1] Processing Tree 0 20%\n",
      "[1] Processing Tree 0 40%\n",
      "[1] Processing Tree 0 60%\n",
      "[1] Processing Tree 0 80%\n",
      "[1] Processing Tree 1 0%\n",
      "[1] Processing Tree 2 0%\n",
      "[1] Processing Tree 3 0%\n",
      "[1] Processing Tree 3 20%\n",
      "[1] Processing Tree 3 40%\n",
      "[1] Processing Tree 3 60%\n",
      "[1] Processing Tree 3 80%\n",
      "[1] Processing Tree 4 0%\n",
      "[1] Processing Tree 4 20%\n",
      "[1] Processing Tree 4 40%\n",
      "[1] Processing Tree 4 60%\n",
      "[1] Processing Tree 4 80%\n",
      "[1] Processing Tree 5 0%\n",
      "[1] Processing Tree 6 0%\n",
      "[1] Processing Tree 7 0%\n",
      "[1] Processing Tree 8 0%\n",
      "[1] Processing Tree 8 20%\n",
      "[1] Processing Tree 8 40%\n",
      "[1] Processing Tree 8 60%\n",
      "[1] Processing Tree 8 80%\n",
      "[1] Processing Tree 9 0%\n",
      "[1] Processing Tree 10 0%\n",
      "[1] Processing Tree 11 0%\n",
      "[1] Processing Tree 12 0%\n",
      "[1] Processing Tree 13 0%\n",
      "[1] Processing Tree 13 20%\n",
      "[1] Processing Tree 13 40%\n",
      "[1] Processing Tree 13 60%\n",
      "[1] Processing Tree 13 80%\n",
      "[1] Processing Tree 14 0%\n",
      "[1] Processing Tree 15 0%\n",
      "[1] Processing Tree 16 0%\n",
      "[1] Processing Tree 17 0%\n",
      "[1] Processing Tree 18 0%\n",
      "[1] Processing Tree 19 0%\n",
      "[1] Processing Tree 19 20%\n",
      "[1] Processing Tree 19 40%\n",
      "[1] Processing Tree 19 60%\n",
      "[1] Processing Tree 19 80%\n",
      "[1] Processing Tree 20 0%\n",
      "[1] Processing Tree 21 0%\n",
      "[1] Processing Tree 22 0%\n",
      "[1] Processing Tree 23 0%\n",
      "[1] Processing Tree 24 0%\n",
      "[1] Processing Tree 25 0%\n",
      "[1] Processing Tree 26 0%\n",
      "[1] Processing Tree 27 0%\n",
      "[1] Processing Tree 28 0%\n",
      "[1] Processing Tree 28 20%\n",
      "[1] Processing Tree 28 40%\n",
      "[1] Processing Tree 28 60%\n",
      "[1] Processing Tree 28 80%\n",
      "[1] Processing Tree 29 0%\n",
      "[1] Processing Tree 30 0%\n",
      "[1] Processing Tree 31 0%\n",
      "[1] Process TPTree time (s): 485\n",
      "[1] BuildInitKNNGraph time (s): 517\n",
      "[1] Refine 0 0%\n",
      "[4] Hash table is full! Set HashTableExponent to larger value (default is 2). NewHashTableExponent=3 NewPoolSize=131071\n",
      "[4] Hash table is full! Set HashTableExponent to larger value (default is 2). NewHashTableExponent=3 NewPoolSize=131071\n",
      "[4] Hash table is full! Set HashTableExponent to larger value (default is 2). NewHashTableExponent=3 NewPoolSize=131071\n",
      "[4] Hash table is full! Set HashTableExponent to larger value (default is 2). NewHashTableExponent=3 NewPoolSize=131071\n",
      "[1] Refine 0 20%\n",
      "[1] Refine 0 40%\n",
      "[1] Refine 0 60%\n",
      "[1] Refine 0 80%\n",
      "[1] Refine RNG time (s): 2080 Graph Acc: 0.997656\n",
      "[1] Refine 1 0%\n",
      "[1] Refine 1 20%\n",
      "[1] Refine 1 40%\n",
      "[1] Refine 1 60%\n",
      "[1] Refine 1 80%\n",
      "[1] Refine RNG time (s): 928 Graph Acc: 0.997187\n",
      "[1] BuildGraph time (s): 3530\n",
      "[1] Build Graph time (s): 3530\n",
      "[1] Save Vector (1000000,128) Finish!\n",
      "[1] Save KDT (1,1000000) Finish!\n",
      "[1] Save RNG (1000000,32) Finish!\n",
      "[1] Save DeleteID (1000000,1) Finish!\n"
     ]
    }
   ],
   "source": [
    "import SPTAG\n",
    "import numpy as np\n",
    "\n",
    "# 假设 SIFT1M fbin 文件路径\n",
    "fbin_path = '/home/gary/Code/DiskANN/build/data/sift/sift_base.fbin'\n",
    "\n",
    "# 读取fbin文件\n",
    "def read_fbin(file_path):\n",
    "    with open(file_path, 'rb') as f:\n",
    "        num_vectors = np.fromfile(f, dtype=np.int32, count=1)[0]  # 读取向量数量\n",
    "        dim = np.fromfile(f, dtype=np.int32, count=1)[0]  # 读取维度\n",
    "        print(\"num_vectors:\",num_vectors,\"dim:\",dim)\n",
    "        print(\"Reading...\")\n",
    "        data = np.fromfile(f, dtype=np.float32).reshape(num_vectors, dim)  # 读取所有向量并重塑为矩阵\n",
    "        print(\"Read done\")\n",
    "    return data\n",
    "\n",
    "# 加载SIFT1M数据\n",
    "sift1m_data = read_fbin(fbin_path)\n",
    "\n",
    "# 创建SPTAG索引实例\n",
    "index_builder = SPTAG.AnnIndex('KDT', 'Float', sift1m_data.shape[1])\n",
    "index_builder.SetBuildParam('NumberOfThreads', '4', \"Index\")  # 设置构建时使用的线程数\n",
    "index_builder.SetBuildParam('Samples', '1000', \"Index\")       # 设置采样数量\n",
    "\n",
    "# 插入数据\n",
    "print(\"Training...\")\n",
    "index_builder.Build(sift1m_data, sift1m_data.shape[0], False)\n",
    "print(\"Training Done...\")\n",
    "\n",
    "# 保存索引\n",
    "index_builder.Save('sift1m_sptag_index2')\n",
    "\n",
    "print(\"索引构建完成并已保存。\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
